/* File: ast_stmt.cc
 * -----------------
 * Implementation of statement node classes.
 */
#include "ast_stmt.h"
#include "ast_type.h"
#include "ast_decl.h"
#include "ast_expr.h"
#include "hashtable.h"
#include "codegen.h"

Scope::Scope() : map(new Hashtable<Decl*>){}

void Scope::insertDecl(Decl *d){
        map->Enter(d->getName(),d);
}

stack<const char*> *Program::labels = new stack<const char*>;
Scope *Program::globalScope = new Scope;

Program::Program(List<Decl*> *d) : generator(new CodeGenerator) {
    Assert(d != NULL);
    this->scope = globalScope;
    (decls=d)->SetParentAll(this);
}

void Program::Check() {
        for( int i = 0; i < decls->NumElements(); i++ ) {
                globalScope->insertDecl(decls->Nth(i));
        }
        for( int i = 0; i < decls->NumElements(); i++ ) {
                decls->Nth(i)->useScope();
        }
}
void Program::Emit() {
        int o = CodeGenerator::OffsetToFirstGlobal;
        for (int i = 0; i < decls->NumElements(); i++) {
                VarDecl *d = dynamic_cast<VarDecl*>(decls->Nth(i));
                if ( d == NULL ) {
                        continue;
                }
                Location *l = new Location(gpRelative, o, d->getName());
                d->setLocation(l);
                o = o + d->getSpace();
        }
        for (int i = 0; i < decls->NumElements(); i++) {
                decls->Nth(i)->locationSet();
        }
        for (int i = 0; i < decls->NumElements(); i++) {
                decls->Nth(i)->codeGen(generator);
                
        }
        generator->DoFinalCodeGen();
}

Stmt::Stmt() : Node() {
        scope = new Scope;
}

Stmt::Stmt(yyltype loc) : Node(loc) {
        scope = new Scope;
}

StmtBlock::StmtBlock(List<VarDecl*> *d, List<Stmt*> *s) {
    Assert(d != NULL && s != NULL);
    (decls=d)->SetParentAll(this);
    (stmts=s)->SetParentAll(this);
}

Location *StmtBlock::codeGen(CodeGenerator *gen) {
        for (int i = 0; i < decls->NumElements(); i++) {
                VarDecl *varDecl = dynamic_cast<VarDecl*>(decls->Nth(i));
                if ( varDecl != NULL ) {
                        Location *b = gen->GenLocalVar(varDecl->getName(),varDecl->getSpace());
                        varDecl->setLocation(b);
                }
        }
        for (int i = 0; i < stmts->NumElements(); i++) {
                stmts->Nth(i)->codeGen(gen);
        }
        return NULL;
} 

void StmtBlock::useScope() {
        for(int i = 0; i < decls->NumElements(); i++) {
                scope->insertDecl(decls->Nth(i));
        }
        for(int i = 0; i < decls->NumElements(); i++) {
                decls->Nth(i)->useScope();
        }
        for(int i = 0; i < stmts->NumElements(); i++){
                stmts->Nth(i)->useScope();
        }
}

int StmtBlock::getSpace() {
        int b = 0;
        for(int i = 0; i < decls->NumElements(); i++) {
                b = b + decls->Nth(i)->getSpace();
        }
        for(int i = 0; i < stmts->NumElements(); i++) {
                b = b + stmts->Nth(i)->getSpace();
        }
        return b;

}

ConditionalStmt::ConditionalStmt(Expr *t, Stmt *b) { 
    Assert(t != NULL && b != NULL);
    (test=t)->SetParent(this); 
    (body=b)->SetParent(this);
}

void ConditionalStmt::useScope() {
        test->useScope();
        body->useScope();
}

ForStmt::ForStmt(Expr *i, Expr *t, Expr *s, Stmt *b): LoopStmt(t, b) { 
    Assert(i != NULL && t != NULL && s != NULL && b != NULL);
    (init=i)->SetParent(this);
    (step=s)->SetParent(this);
}

void ForStmt::useScope() {
        LoopStmt::useScope();
        init->useScope();
        step->useScope();
}

int ForStmt::getSpace() {
        return step->getSpace() + init->getSpace() + test->getSpace() + body->getSpace();
}

Location* ForStmt::codeGen(CodeGenerator *gen) {
        const char* end = gen->NewLabel();
        Program::labels->push(end);
        init->codeGen(gen);
        const char* beg = gen->NewLabel();
        gen ->GenLabel(beg);
        Location *l = test->codeGen(gen);
        gen->GenIfZ(l,end);
        body->codeGen(gen);
        step->codeGen(gen);
        gen->GenGoto(beg);
        gen->GenLabel(end);
        Program::labels->pop();
        return NULL;
}

int WhileStmt::getSpace() {
        return body->getSpace() + test->getSpace();
}
void WhileStmt::useScope() {
        LoopStmt::useScope();
}

Location *WhileStmt::codeGen(CodeGenerator *gen) {
        const char* end = gen->NewLabel();
        Program::labels->push(end);
        const char* start = gen->NewLabel();
        gen->GenLabel(start);
        Location *l = test->codeGen(gen);
        gen->GenIfZ(l, end);
        body->codeGen(gen);
        gen->GenGoto(start);
        gen->GenLabel(end);
        Program::labels->pop();
        return NULL;
}

IfStmt::IfStmt(Expr *t, Stmt *tb, Stmt *eb): ConditionalStmt(t, tb) { 
    Assert(t != NULL && tb != NULL); // else can be NULL
    elseBody = eb;
    if (elseBody) elseBody->SetParent(this);
}

int IfStmt::getSpace() {
        int b = body->getSpace() + test->getSpace();
        if ( elseBody != NULL ) {
                b = b + elseBody->getSpace();
        }
        return b;
}

void LoopStmt::useScope() {
        ConditionalStmt::useScope();
}

Location *IfStmt::codeGen(CodeGenerator *gen) {
        const char* end = gen->NewLabel();
        const char* elseStmt = gen->NewLabel();
        Location *l = test->codeGen(gen);
        gen->GenIfZ(l,elseStmt);
        body->codeGen(gen);
        gen->GenGoto(end);
        gen->GenLabel(elseStmt);
        if ( elseBody != NULL ) {
                elseBody->codeGen(gen);
        }
        gen->GenLabel(end);
        return NULL;
}

void IfStmt::useScope() {
        ConditionalStmt::useScope();
        if (elseBody != NULL) {
                elseBody->useScope();
        }
}

ReturnStmt::ReturnStmt(yyltype loc, Expr *e) : Stmt(loc) { 
    Assert(e != NULL);
    (expr=e)->SetParent(this);
}

Location *ReturnStmt::codeGen(CodeGenerator *gen) {
        if (expr != NULL) {
                gen->GenReturn(expr->codeGen(gen));
        } else {
                gen->GenReturn();
        }
        return NULL;
}

int ReturnStmt::getSpace() {
        if ( expr != NULL ) {
                return expr->getSpace();
        } else {
                return 0;
        }
}
  
void ReturnStmt::useScope() { expr->useScope(); }

  
  
PrintStmt::PrintStmt(List<Expr*> *a) {    
    Assert(a != NULL);
    (args=a)->SetParentAll(this);
}

void PrintStmt::useScope() {
        for ( int i = 0; i < args->NumElements(); i++) {
                args->Nth(i)->useScope();
        }
}

int PrintStmt::getSpace() {
        int b = 0;
        for (int i = 0; i < args->NumElements(); i++) {
                b = b + args->Nth(i)->getSpace();
        }
        return b;
}

Location *PrintStmt::codeGen(CodeGenerator *gen) {
        for (int i = 0; i < args->NumElements(); i++) {
                BuiltIn print = args->Nth(i)->getType()->getPrint();
                Assert(print != NumBuiltIns);
                gen->GenBuiltInCall(print, args->Nth(i)->codeGen(gen));
        }
        return NULL;
}

int BreakStmt::getSpace() { return 0; }

Location *BreakStmt::codeGen(CodeGenerator *gen) {
        gen->GenGoto(Program::labels->top());
        return NULL;
}